import argparse
import logging
import math
import os
from pathlib import Path

import accelerate
import numpy as np
import torch
import torch.nn.functional as F
import torch.utils.checkpoint
import torchvision
from accelerate import Accelerator
from accelerate.logging import get_logger
from accelerate.utils import ProjectConfiguration, set_seed
from datasets import load_dataset
from huggingface_hub import create_repo, upload_folder
from packaging import version
from torchvision import transforms
from tqdm.auto import tqdm

import diffusers
# from diffusers import AutoencoderKL
from autoencoder_kl import AutoencoderKL
from diffusers.optimization import get_scheduler
from diffusers.training_utils import EMAModel
from diffusers.utils import is_wandb_available

from discriminator_vae import Discriminator
import lpips
from PIL import Image

from PIL import PngImagePlugin
PngImagePlugin.MAX_TEXT_CHUNK = 1000 * (1024**2)
Image.MAX_IMAGE_PIXELS = None

from PIL import ImageFile
ImageFile.LOAD_TRUNCATED_IMAGES = True

if is_wandb_available():
    import wandb

logger = get_logger(__name__, log_level="INFO")


def grad_layer_wrt_loss(loss, layer):
    return torch.autograd.grad(
        outputs=loss,
        inputs=layer,
        grad_outputs=torch.ones_like(loss),
        retain_graph=True,
    )[0].detach()

def gradient_penalty(images, output, weight=10):
    gradients = torch.autograd.grad(
        outputs=output,
        inputs=images,
        grad_outputs=torch.ones(output.size(), device=images.device),
        create_graph=True,
        retain_graph=True,
        only_inputs=True,
    )[0]
    bsz = gradients.shape[0]
    gradients = torch.reshape(gradients, (bsz, -1))
    return weight * ((gradients.norm(2, dim=1) - 1) ** 2).mean()

def log_grad_norm(model, accelerator, global_step):
    for name, param in model.named_parameters():
        if param.grad is not None:
            grads = param.grad.detach().data
            grad_norm = (grads.norm(p=2) / grads.numel()).item()
            accelerator.log({"grad_norm/" + name: grad_norm}, step=global_step)

# Function to split the image into patches
def extract_patches(image, patch_size, stride):
    # Unfold the image into patches
    patches = image.unfold(2, patch_size, stride).unfold(3, patch_size, stride)
    # Reshape to get a batch of patches
    patches = patches.contiguous().view(image.size(0), image.size(1), -1, patch_size, patch_size)
    return patches

# Patch-Based MSE Loss
def patch_based_mse_loss(real_images, recon_images, patch_size=32, stride=16):
    real_patches = extract_patches(real_images, patch_size, stride)
    recon_patches = extract_patches(recon_images, patch_size, stride)
    mse_loss = F.mse_loss(real_patches, recon_patches)
    return mse_loss

# Patch-Based LPIPS Loss (using the pre-defined LPIPS model)
def patch_based_lpips_loss(lpips_model, real_images, recon_images, patch_size=32, stride=16):
    with torch.no_grad():
        real_patches = extract_patches(real_images, patch_size, stride)
        recon_patches = extract_patches(recon_images, patch_size, stride)
        
        lpips_loss = 0
        # Iterate over each patch and accumulate LPIPS loss
        for i in range(real_patches.size(2)):  # Loop over number of patches
            real_patch = real_patches[:, :, i, :, :].contiguous()
            recon_patch = recon_patches[:, :, i, :, :].contiguous()
            patch_lpips_loss = lpips_model(real_patch, recon_patch).mean()
            
            # Handle non-finite values
            if not torch.isfinite(patch_lpips_loss):
                patch_lpips_loss = torch.tensor(0, device=real_patch.device)
            
            lpips_loss += patch_lpips_loss

    return lpips_loss / real_patches.size(2)  # Normalize by the number of patches


@torch.no_grad()
def log_validation(args, test_dataloader, vae, accelerator, weight_dtype, epoch, prefix="VAE F16"):
    logger.info("Running validation... ")

    vae_model = accelerator.unwrap_model(vae)
    images = []
    
    for _, sample in enumerate(test_dataloader):
        x = sample["pixel_values"].to(weight_dtype)
        # print(x.size())
        reconstructions = vae_model(x).sample
        images.append(
            # torch.cat([sample["pixel_values"].cpu(), reconstructions.cpu()], axis=0)
            torch.cat([sample["pixel_values"].cpu(), reconstructions.cpu()], axis=3)
        )

    for tracker in accelerator.trackers:
        if tracker.name == "tensorboard":
            np_images = np.concatenate([np.asarray((img / 2 + 0.5).clamp(0, 1)) for img in images], axis=0)
            tracker.writer.add_images(
                "{}, Original (left), Reconstruction (right), epoch {}".format(prefix, epoch), np_images, epoch
            )

            # for idx, image in enumerate(images):
            #     tracker.writer.add_images(
            #         "Original (left), Reconstruction (right), image {}, epoch {}".format(idx, epoch), np.asarray((image / 2 + 0.5).clamp(0, 1)), epoch
            #     )

        elif tracker.name == "wandb":
            tracker.log(
                {
                    "Original (left), Reconstruction (right)": [
                        wandb.Image(torchvision.utils.make_grid(image))
                        for _, image in enumerate(images)
                    ]
                }
            )
        else:
            logger.warn(f"image logging not implemented for {tracker.gen_images}")

    del vae_model
    torch.cuda.empty_cache()

def make_image_grid(imgs, rows, cols):

    w, h = imgs[0].size
    grid = Image.new("RGB", size=(cols * w, rows * h))

    for i, img in enumerate(imgs):
        grid.paste(img, box=(i % cols * w, i // cols * h))
    return grid

def parse_args():
    parser = argparse.ArgumentParser(
        description="Simple example of a VAE training script."
    )
    parser.add_argument(
        "--pretrained_model_name_or_path",
        type=str,
        default=None,
        required=False,
        help="Path to pretrained model or model identifier from huggingface.co/models.",
    )
    parser.add_argument(
        "--discriminator_pretrained_model_name_or_path",
        type=str,
        default=None,
        help="The config of the discriminator model to train, leave as None to use standard Vq model configuration.",
    )
    parser.add_argument(
        "--revision",
        type=str,
        default=None,
        required=False,
        help="Revision of pretrained model identifier from huggingface.co/models.",
    )
    parser.add_argument(
        "--dataset_name",
        type=str,
        default=None,
        help=(
            "The name of the Dataset (from the HuggingFace hub) to train on (could be your own, possibly private,"
            " dataset). It can also be a path pointing to a local copy of a dataset in your filesystem,"
            " or to a folder containing files that 🤗 Datasets can understand."
        ),
    )
    parser.add_argument(
        "--dataset_config_name",
        type=str,
        default=None,
        help="The config of the Dataset, leave as None if there's only one config.",
    )
    parser.add_argument(
        "--train_data_dir",
        type=str,
        default=None,
        help=(
            "A folder containing the training data. Folder contents must follow the structure described in"
            " https://huggingface.co/docs/datasets/image_dataset#imagefolder. In particular, a `metadata.jsonl` file"
            " must exist to provide the captions for the images. Ignored if `dataset_name` is specified."
        ),
    )
    parser.add_argument(
        "--test_data_dir",
        type=str,
        default=None,
        help=(
            "A folder containing the validation data. Folder contents must follow the structure described in"
            " https://huggingface.co/docs/datasets/image_dataset#imagefolder. In particular, a `metadata.jsonl` file"
            " must exist to provide the captions for the images. Ignored if `dataset_name` is specified."
        ),
    )
    parser.add_argument(
        "--image_column",
        type=str,
        default="image",
        help="The column of the dataset containing an image.",
    )
    parser.add_argument(
        "--output_dir",
        type=str,
        default="outputs",
        help="The output directory where the model predictions and checkpoints will be written.",
    )
    parser.add_argument(
        "--huggingface_repo",
        type=str,
        default="vae-model-finetuned",
        help="The output directory where the model predictions and checkpoints will be written.",
    )
    parser.add_argument(
        "--cache_dir",
        type=str,
        default=None,
        help="The directory where the downloaded models and datasets will be stored.",
    )
    parser.add_argument(
        "--seed", type=int, default=21, help="A seed for reproducible training."
    )
    parser.add_argument(
        "--resolution",
        type=int,
        default=512,#512,
        help=(
            "The resolution for input images, all the images in the train/validation dataset will be resized to this"
            " resolution"
        ),
    )
    parser.add_argument(
        "--eval_resolution",
        type=int,
        default=1024,#512,
        help=(
            "The resolution for input images, all the images in the train/validation dataset will be resized to this"
            " resolution"
        ),
    )
    parser.add_argument(
        "--train_batch_size",
        type=int,
        default=1,
        help="Batch size (per device) for the training dataloader.",
    )
    parser.add_argument("--num_train_epochs", type=int, default=100)
    parser.add_argument(
        "--gradient_accumulation_steps",
        type=int,
        default=2,
        help="Number of updates steps to accumulate before performing a backward/update pass.",
    )
    parser.add_argument(
        "--gradient_checkpointing",
        action="store_true",
        help="Whether or not to use gradient checkpointing to save memory at the expense of slower backward pass.",
    )
    parser.add_argument(
        "--learning_rate",
        type=float,
        default=1.5e-7, # Reference : Waifu-diffusion-v1-4 config
        # default=4.5e-8,
        help="Initial learning rate (after the potential warmup period) to use.",
    )
    parser.add_argument(
        "--scale_lr",
        action="store_true",
        default=False,
        help="Scale the learning rate by the number of GPUs, gradient accumulation steps, and batch size.",
    )
    parser.add_argument(
        "--lr_scheduler",
        type=str,
        default="constant",
        help=(
            'The scheduler type to use. Choose between ["linear", "cosine", "cosine_with_restarts", "polynomial",'
            ' "constant", "constant_with_warmup"]'
        ),
    )
    parser.add_argument(
        "--lr_warmup_steps",
        type=int,
        default=500,
        help="Number of steps for the warmup in the lr scheduler.",
    )
    parser.add_argument(
        "--logging_dir",
        type=str,
        default="logs",
        help=(
            "[TensorBoard](https://www.tensorflow.org/tensorboard) log directory. Will default to"
            " *output_dir/runs/**CURRENT_DATETIME_HOSTNAME***."
        ),
    )
    parser.add_argument(
        "--mixed_precision",
        type=str,
        default=None,
        choices=["no", "fp16", "bf16"],
        help=(
            "Whether to use mixed precision. Choose between fp16 and bf16 (bfloat16). Bf16 requires PyTorch >="
            " 1.10.and an Nvidia Ampere GPU.  Default to the value of accelerate config of the current system or the"
            " flag passed with the `accelerate.launch` command. Use this argument to override the accelerate config."
        ),
    )
    parser.add_argument(
        "--report_to",
        type=str,
        default="tensorboard",
        help=(
            'The integration to report the results and logs to. Supported platforms are `"tensorboard"`'
            ' (default), `"wandb"` and `"comet_ml"`. Use `"all"` to report to all integrations.'
        ),
    )
    parser.add_argument(
        "--checkpointing_steps",
        type=int,
        default=5000,
        help=(
            "Save a checkpoint of the training state every X updates. These checkpoints are only suitable for resuming"
            " training using `--resume_from_checkpoint`."
        ),
    )
    parser.add_argument(
        "--checkpoints_total_limit",
        type=int,
        default=None,
        help=(
            "Max number of checkpoints to store. Passed as `total_limit` to the `Accelerator` `ProjectConfiguration`."
            " See Accelerator::save_state https://huggingface.co/docs/accelerate/package_reference/accelerator#accelerate.Accelerator.save_state"
            " for more docs"
        ),
    )
    parser.add_argument(
        "--resume_from_checkpoint",
        type=str,
        default=None,
        help=(
            "Whether training should be resumed from a previous checkpoint. Use a path saved by"
            ' `--checkpointing_steps`, or `"latest"` to automatically select the last available checkpoint.'
        ),
    )
    parser.add_argument(
        "--test_samples",
        type=int,
        default=20,
        help="Number of images to remove from training set to be used as validation.",
    )
    parser.add_argument(
        "--validation_epochs",
        type=int,
        default=1,
        help="Run validation every X epochs.",
    )
    parser.add_argument(
        "--tracker_project_name",
        type=str,
        default="vae-fine-tune",
        help=(
            "The `project_name` argument passed to Accelerator.init_trackers for"
            " more information see https://huggingface.co/docs/accelerate/v0.17.0/en/package_reference/accelerator#accelerate.Accelerator"
        ),
    )
    parser.add_argument(
        "--use_8bit_adam", action="store_true", help="Whether or not to use 8-bit Adam from bitsandbytes."
    )
    parser.add_argument("--use_ema", action="store_true", help="Whether to use EMA model.")
    parser.add_argument(
        "--kl_scale",
        type=float,
        default=1e-6,
        help="Scaling factor for the Kullback-Leibler divergence penalty term.",
    )
    parser.add_argument(
        "--lpips_scale",
        type=float,
        default=5e-1,
        help="Scaling factor for the LPIPS metric",
    )
    parser.add_argument(
        "--lpips_start",
        type=int,
        default=50001,
        help="Start for the LPIPS metric",
    )
    parser.add_argument("--adam_beta1", type=float, default=0.9, help="The beta1 parameter for the Adam optimizer.")
    parser.add_argument("--adam_beta2", type=float, default=0.999, help="The beta2 parameter for the Adam optimizer.")
    parser.add_argument("--adam_weight_decay", type=float, default=1e-2, help="Weight decay to use.")
    parser.add_argument("--adam_epsilon", type=float, default=1e-08, help="Epsilon value for the Adam optimizer")
    parser.add_argument("--max_grad_norm", default=1.0, type=float, help="Max gradient norm.")
    parser.add_argument("--push_to_hub", action="store_true", help="Whether or not to push the model to the Hub.")
    parser.add_argument("--hub_token", type=str, default=None, help="The token to use to push to the Model Hub.")
    parser.add_argument("--patch_loss", action="store_true", help="Whether to use patch loss.")

    parser.add_argument(
        "--discr_learning_rate",
        type=float,
        default=1e-4,
        help="Initial learning rate (after the potential warmup period) to use.",
    )
    parser.add_argument(
        "--discr_lr_scheduler",
        type=str,
        default="constant",
        help=(
            'The scheduler type to use. Choose between ["linear", "cosine", "cosine_with_restarts", "polynomial",'
            ' "constant", "constant_with_warmup"]'
        ),
    )
    parser.add_argument(
        "--gen_scale",
        type=float,
        default=5e-2,
        help="Scaling factor for the gan penalty term.",
    )
    
    parser.add_argument(
        "--sc_scale",
        type=float,
        default=1,
        help="Scaling factor for the scale consistency term.",
    )
    
    args = parser.parse_args()
    # Sanity checks
    if args.dataset_name is None and args.train_data_dir is None:
        raise ValueError("Need either a dataset name or a training folder.")

    return args


def main():
    args = parse_args()

    logging_dir = os.path.join(args.output_dir, args.logging_dir)

    accelerator_project_config = ProjectConfiguration(
        total_limit=args.checkpoints_total_limit,
        project_dir=args.output_dir,
        logging_dir=logging_dir,
    )

    accelerator = Accelerator(
        gradient_accumulation_steps=args.gradient_accumulation_steps,
        mixed_precision=args.mixed_precision,
        log_with=args.report_to,
        project_config=accelerator_project_config,
    )

    # Make one log on every process with the configuration for debugging.
    logging.basicConfig(
        format="%(asctime)s - %(levelname)s - %(name)s - %(message)s",
        datefmt="%m/%d/%Y %H:%M:%S",
        level=logging.INFO,
    )
    logger.info(accelerator.state, main_process_only=False)

    if args.seed is not None:
        set_seed(args.seed)

    if accelerator.is_main_process:
        if args.output_dir is not None:
            os.makedirs(args.output_dir, exist_ok=True)
            

    weight_dtype = torch.float32
    if accelerator.mixed_precision == "fp16":
        weight_dtype = torch.float16
    elif accelerator.mixed_precision == "bf16":
        weight_dtype = torch.bfloat16

    # Load vae
    try:
        vae = AutoencoderKL.from_pretrained(
            args.pretrained_model_name_or_path, subfolder="vae", revision=args.revision, weight_dtype=torch.float32, stride=2,
        )
        teacher_vae = AutoencoderKL.from_pretrained(
            args.pretrained_model_name_or_path, subfolder="vae", revision=args.revision, weight_dtype=torch.float32, stride=1,
        )
    except:
        vae = AutoencoderKL.from_pretrained(
            args.pretrained_model_name_or_path, revision=args.revision, weight_dtype=torch.float32, stride=2,
        )
        teacher_vae = AutoencoderKL.from_pretrained(
            args.pretrained_model_name_or_path, revision=args.revision, weight_dtype=torch.float32, stride=1,
        )
    if args.use_ema:
        try:
            ema_vae = AutoencoderKL.from_pretrained(
                args.pretrained_model_name_or_path, subfolder="vae", revision=args.revision, weight_dtype=torch.float32, stride=2,)
        except:
            ema_vae = AutoencoderKL.from_pretrained(
                args.pretrained_model_name_or_path, revision=args.revision, weight_dtype=torch.float32, stride=2)
        ema_vae = EMAModel(ema_vae.parameters(), model_cls=AutoencoderKL, model_config=ema_vae.config)

    if args.discriminator_pretrained_model_name_or_path is None:
        discriminator = Discriminator()
    else:
        discriminator = Discriminator.from_pretrained(args.discriminator_pretrained_model_name_or_path)


    teacher_vae.requires_grad_(False)
    vae.requires_grad_(False)
    vae_trainable_params = []
    for name, param in vae.named_parameters():
        # if 'decoder' in name and 'decoder.conv_out' not in name:
        if 'decoder' in name:
            # print(name)
            param.requires_grad = True
            vae_trainable_params.append(param)

    print(f"VAE total params = {len(list(vae.named_parameters()))}, trainable params = {len(vae_trainable_params)}")

    vae.to(device=accelerator.device, dtype=weight_dtype)
    teacher_vae.to(device=accelerator.device, dtype=weight_dtype)
    discriminator.to(device=accelerator.device, dtype=weight_dtype)

    # `accelerate` 0.16.0 will have better support for customized saving
    if version.parse(accelerate.__version__) >= version.parse("0.16.0"):
        # create custom saving & loading hooks so that `accelerator.save_state(...)` serializes in a nice format
        def save_model_hook(models, weights, output_dir):
            if args.use_ema:
                ema_vae.save_pretrained(os.path.join(output_dir, "vae_ema"))

            # logger.info(f"{vae = }")
            vae = models[0]
            vae.save_pretrained(os.path.join(output_dir, "vae"))
            discriminator = models[1]
            discriminator.save_pretrained(os.path.join(output_dir, "discriminator"))
            weights.pop()
            weights.pop()
            

        def load_model_hook(models, input_dir):
            if args.use_ema:
                load_model = EMAModel.from_pretrained(os.path.join(input_dir, "vae_ema"), AutoencoderKL)
                ema_vae.load_state_dict(load_model.state_dict())
                ema_vae.to(accelerator.device)
                del load_model

            # load diffusers style into model
            vae = models.pop()
            load_model = AutoencoderKL.from_pretrained(input_dir, subfolder="vae")
            vae.register_to_config(**load_model.config)

            vae.load_state_dict(load_model.state_dict())
            del load_model

            discriminator = models.pop()
            load_model = Discriminator.from_pretrained(input_dir, subfolder="discriminator")
            discriminator.register_to_config(**load_model.config)
            discriminator.load_state_dict(load_model.state_dict())
            del load_model

        accelerator.register_save_state_pre_hook(save_model_hook)
        accelerator.register_load_state_pre_hook(load_model_hook)

    if args.gradient_checkpointing:
        vae.enable_gradient_checkpointing()

    if args.scale_lr:
        args.learning_rate = (
            args.learning_rate
            * args.gradient_accumulation_steps
            * args.train_batch_size
            * accelerator.num_processes
        )
        
    # Initialize the optimizer
    if args.use_8bit_adam:
        try:
            import bitsandbytes as bnb
        except ImportError:
            raise ImportError(
                "Please install bitsandbytes to use 8-bit Adam. You can do so by running `pip install bitsandbytes` or `pip install bitsandbytes-windows` for Windows"
            )

        optimizer_cls = bnb.optim.AdamW8bit
    else:
        optimizer_cls = torch.optim.AdamW

    import itertools
    params_to_optimize = (
        itertools.chain(vae_trainable_params)
    )

    params_to_optimize = {"params": params_to_optimize, "lr": args.learning_rate}
    params_to_optimize = [params_to_optimize]

    optimizer = optimizer_cls(
        params_to_optimize,
        # vae.parameters(),
        lr=args.learning_rate,
        betas=(args.adam_beta1, args.adam_beta2),
        weight_decay=args.adam_weight_decay,
        eps=args.adam_epsilon,
    )

    discr_optimizer = optimizer_cls(
        list(discriminator.parameters()),
        lr=args.discr_learning_rate,
        betas=(args.adam_beta1, args.adam_beta2),
        weight_decay=args.adam_weight_decay,
        eps=args.adam_epsilon,
    )

    # Get the datasets: you can either provide your own training and evaluation files (see below)
    # or specify a Dataset from the hub (the dataset will be downloaded automatically from the datasets Hub).

    # In distributed training, the load_dataset function guarantees that only one local process can concurrently
    # download the dataset.
    if args.dataset_name is not None:
        # Downloading and loading a dataset from the hub.
        dataset = load_dataset(
            args.dataset_name,
            args.dataset_config_name,
            cache_dir=args.cache_dir,
        )
    else:
        data_files = {}
        if args.train_data_dir is not None:
            # data_files["train"] = os.path.join(args.train_data_dir, "metadata_vae_train.jsonl")
            data_files["train"] = os.path.join(args.train_data_dir, "metadata_vae_sa1b.jsonl")
            data_files["test"] = os.path.join(args.train_data_dir, "metadata_vae_val.jsonl")

        dataset = load_dataset("json", data_files=data_files) #, split="train")#, streaming=True)

    column_names = dataset["train"].column_names
    if args.image_column is None:
        image_column = column_names[0]
    else:
        image_column = args.image_column
        if image_column not in column_names:
            raise ValueError(
                f"--image_column' value '{args.image_column}' needs to be one of: {', '.join(column_names)}"
            )

    train_transforms = transforms.Compose(
        [
            transforms.Resize(
                args.resolution, interpolation=transforms.InterpolationMode.BILINEAR
            ),
            transforms.RandomCrop(args.resolution),
            # transforms.RandomResizedCrop(args.resolution, scale=(0.6, 1.0)),
            # transforms.RandomApply(
            #     [transforms.ColorJitter(0.2, 0.2, 0.2, 0.1)],
            #     p=0.5,  # not strengthened
            # ),
            transforms.RandomHorizontalFlip(),
            transforms.ToTensor(),
            transforms.Normalize([0.5], [0.5]),
        ]
    )

    test_transforms = transforms.Compose(
        [
            transforms.Resize(
                args.eval_resolution, interpolation=transforms.InterpolationMode.BILINEAR
            ),
            transforms.CenterCrop(args.eval_resolution),
            transforms.ToTensor(),
            transforms.Normalize([0.5], [0.5]),
        ]
    )
    

    def preprocess(examples):
        images = [Image.open(os.path.join(args.train_data_dir, image_path)).convert("RGB") for image_path in examples[image_column]]
        examples["pixel_values"] = [train_transforms(image) for image in images]
        return examples
    
    def test_preprocess(examples):
        images = [Image.open(os.path.join(args.train_data_dir, image_path)).convert("RGB") for image_path in examples[image_column]]
        examples["pixel_values"] = [test_transforms(image) for image in images]
        return examples

    with accelerator.main_process_first():
        # Load test data from test_data_dir
        if(args.test_data_dir is not None and args.train_data_dir is not None):
            logger.info(f"load test data from {args.test_data_dir}")
            test_dir = os.path.join(args.test_data_dir, "**")        
            test_dataset = load_dataset(
                "imagefolder",
                data_files=test_dir,
                cache_dir=args.cache_dir,
            )
            # Set the training transforms
            train_dataset = dataset["train"].with_transform(preprocess)
            test_dataset = test_dataset["train"].with_transform(test_preprocess)
        # Load train/test data from train_data_dir
        elif "test" in dataset.keys():
            train_dataset = dataset["train"].with_transform(preprocess)
            test_dataset = dataset["test"].with_transform(test_preprocess)
        # Split into train/test
        else:
            dataset = dataset["train"].train_test_split(test_size=args.test_samples)        
            # Set the training transforms
            train_dataset = dataset["train"].with_transform(preprocess)
            test_dataset = dataset["test"].with_transform(test_preprocess)

    def collate_fn(examples):
        pixel_values = torch.stack([example["pixel_values"] for example in examples])
        pixel_values = pixel_values.to(memory_format=torch.contiguous_format).float()
        return {"pixel_values": pixel_values}

    # DataLoaders creation:
    train_dataloader = torch.utils.data.DataLoader(
        train_dataset,
        shuffle=True,
        collate_fn=collate_fn,
        batch_size=args.train_batch_size,
        num_workers=4,
        # num_workers=args.train_batch_size*accelerator.num_processes,
    )

    test_dataloader = torch.utils.data.DataLoader(
        test_dataset, shuffle=False, collate_fn=collate_fn, batch_size=16, num_workers=2,#args.train_batch_size*accelerator.num_processes,
    )

    lr_scheduler = get_scheduler(
        args.lr_scheduler,
        optimizer=optimizer,
        num_warmup_steps=args.lr_warmup_steps * args.gradient_accumulation_steps,
        num_training_steps=args.num_train_epochs * args.gradient_accumulation_steps,
    )

    discr_lr_scheduler = get_scheduler(
        args.discr_lr_scheduler,
        optimizer=discr_optimizer,
        num_warmup_steps=args.lr_warmup_steps * args.gradient_accumulation_steps,
        num_training_steps=args.num_train_epochs * args.gradient_accumulation_steps,
    )

    
    (vae,
        discriminator,
        teacher_vae,
        optimizer,
        discr_optimizer,
        train_dataloader,
        test_dataloader,
        lr_scheduler,
        discr_lr_scheduler,
    ) = accelerator.prepare(
        vae, discriminator, teacher_vae, optimizer, discr_optimizer, train_dataloader, test_dataloader, lr_scheduler, discr_lr_scheduler
    )


    if args.use_ema:
        ema_vae.to(device=accelerator.device, dtype=weight_dtype) 
    
    # We need to initialize the trackers we use, and also store our configuration.
    # The trackers initializes automatically on the main process.
    if accelerator.is_main_process:
        tracker_config = dict(vars(args))
        accelerator.init_trackers(args.tracker_project_name, tracker_config)

    num_update_steps_per_epoch = math.ceil(
        len(train_dataloader) / args.gradient_accumulation_steps
    )
    args.max_train_steps = args.num_train_epochs * num_update_steps_per_epoch

    # ------------------------------ TRAIN ------------------------------ #
    total_batch_size = (
        args.train_batch_size
        * accelerator.num_processes
        * args.gradient_accumulation_steps
    )

    logger.info("***** Running training *****")
    logger.info(f"  Num examples = {len(train_dataset)}")
    logger.info(f"  Num test samples = {len(test_dataset)}")
    logger.info(f"  Num Epochs = {args.num_train_epochs}")
    logger.info(f"  Instantaneous batch size per device = {args.train_batch_size}")
    logger.info(
        f"  Total train batch size (w. parallel, distributed & accumulation) = {total_batch_size}"
    )
    logger.info(f"  Gradient Accumulation steps = {args.gradient_accumulation_steps}")
    global_step = 0
    first_epoch = 0

    # Potentially load in the weights and states from a previous save
    if args.resume_from_checkpoint:
        if args.resume_from_checkpoint != "latest":
            path = os.path.basename(args.resume_from_checkpoint)
        else:
            # Get the most recent checkpoint
            dirs = os.listdir(args.output_dir)
            # dirs = os.listdir(args.resume_from_checkpoint)
            dirs = [d for d in dirs if d.startswith("checkpoint")]
            dirs = sorted(dirs, key=lambda x: int(x.split("-")[1]))
            path = dirs[-1] if len(dirs) > 0 else None

        if path is None:
            print(
                f"Checkpoint '{args.resume_from_checkpoint}' does not exist. Starting a new training run."
            )
            args.resume_from_checkpoint = None
        else:
            print(f"Resuming from checkpoint {path}")
            accelerator.load_state(os.path.join(args.output_dir, path))
            # accelerator.load_state(os.path.join(path)) #kiml
            global_step = int(path.split("-")[1])

            # resume_global_step = global_step * args.gradient_accumulation_steps
            # first_epoch = global_step // num_update_steps_per_epoch
            # resume_step = resume_global_step % (num_update_steps_per_epoch * args.gradient_accumulation_steps)

            # resume_global_step = global_step #* args.gradient_accumulation_steps
            first_epoch = global_step // num_update_steps_per_epoch
            resume_step = global_step # resume_global_step % (num_update_steps_per_epoch * args.gradient_accumulation_steps)

    progress_bar = tqdm(
        range(0, args.max_train_steps),
        initial=global_step,
        disable=not accelerator.is_local_main_process,
    )
    progress_bar.set_description("Steps")

    # lpips_loss_fn = lpips.LPIPS(net="alex").to(accelerator.device, dtype=weight_dtype)
    lpips_loss_fn = lpips.LPIPS(net="vgg").to(accelerator.device, dtype=weight_dtype)
    lpips_loss_fn.requires_grad_(False)
    # lpips_loss_fn.eval()

    for epoch in range(first_epoch, args.num_train_epochs):
        vae.train()
        discriminator.train()
        teacher_vae.eval()
        logger.info(f"{epoch = }")

        for step, batch in enumerate(train_dataloader):
            # # Skip steps until we reach the resumed step
            # if args.resume_from_checkpoint and epoch == first_epoch and step < resume_step:
            #     if step % args.gradient_accumulation_steps == 0:
            #         progress_bar.update(1)
            #     continue
            
            target = batch["pixel_values"].to(weight_dtype)
            generator_step = ((step // args.gradient_accumulation_steps) % 2) == 0

            # Train Step
            # The behavior of accelerator.accumulate is to
            # 1. Check if gradients are synced(reached gradient-accumulation_steps)
            # 2. If so sync gradients by stopping the not syncing process
            if generator_step:
                optimizer.zero_grad(set_to_none=True)
            else:
                discr_optimizer.zero_grad(set_to_none=True)

            if generator_step:
                with accelerator.accumulate(vae):
                    # https://github.com/huggingface/diffusers/blob/main/src/diffusers/models/autoencoder_kl.py
                    if accelerator.num_processes > 1:
                        posterior = vae.module.encode(target).latent_dist
                        # teacher_posterior = teacher_vae.module.encode(target).latent_dist
                        teacher_posterior = teacher_vae.encode(target).latent_dist
                    else:
                        posterior = accelerator.unwrap_model(vae).encode(target).latent_dist
                        teacher_posterior = accelerator.unwrap_model(teacher_vae).encode(target).latent_dist
                        
                    
                    # z = mean                      if posterior.mode()
                    # z = mean + variable*epsilon   if posterior.sample()
                    z = posterior.sample() # Not mode()
                    teacher_z = teacher_posterior.sample()

                    # z = z * 1.5305  # * vae.config.scaling_factor  sd3: 1.5305   flux: 0.3611
                    # z = 1 / 1.5305 * z

                    partitioned = True # if np.random.rand() > 0.5 else False

                    if accelerator.num_processes > 1:
                        pred = vae.module.decode(z, partitioned=partitioned).sample
                        immediate_layer = vae.module.decode(z, partitioned=partitioned).immediate_layer
                        teacher_immediate_layer = teacher_vae.decode(teacher_z, partitioned=partitioned).immediate_layer
                    else:
                        pred = accelerator.unwrap_model(vae).decode(z, partitioned=partitioned).sample
                        immediate_layer = accelerator.unwrap_model(vae).decode(z, partitioned=partitioned).immediate_layer
                        teacher_immediate_layer = accelerator.unwrap_model(teacher_vae).decode(teacher_z, partitioned=partitioned).immediate_layer

                    pred = pred.clamp(-1, 1)

                    # remove kl loss to train decoder
                    # kl_loss = posterior.kl().mean()

                    if args.patch_loss:
                        # patched loss
                        mse_loss = patch_based_mse_loss(target, pred, patch_size=32, stride=16).to(weight_dtype)
                        lpips_loss = patch_based_lpips_loss(lpips_loss_fn, target, pred, patch_size=32, stride=16).to(weight_dtype)

                    else:
                        mse_loss = F.mse_loss(pred.float(), target.float(), reduction="mean")  
                    
                        # with torch.no_grad():
                        lpips_loss = lpips_loss_fn(pred.to(dtype=weight_dtype), target).mean()
                        if not torch.isfinite(lpips_loss):
                            lpips_loss = torch.tensor(0)

                    # generator loss
                    gen_loss = -discriminator(pred).mean()

                    # last_dec_layer = accelerator.unwrap_model(vae).decoder.conv_out.gate.weight
                    # last_dec_layer = accelerator.unwrap_model(vae).decoder.conv_norm_out.weight
                    last_dec_layer = accelerator.unwrap_model(vae).decoder.conv_out.weight
                    norm_grad_wrt_perceptual_loss = grad_layer_wrt_loss(lpips_loss, last_dec_layer).norm(p=2)
                    norm_grad_wrt_gen_loss = grad_layer_wrt_loss(gen_loss, last_dec_layer).norm(p=2)

                    adaptive_weight = norm_grad_wrt_perceptual_loss / norm_grad_wrt_gen_loss.clamp(min=1e-8)
                    adaptive_weight = adaptive_weight.clamp(max=1e4)


                    # sc_loss = F.mse_loss(F.interpolate(teacher_immediate_layer.detach().float(), immediate_layer.shape[2:], mode='nearest'), immediate_layer.float(), reduction="mean") 
                    sc_loss = F.mse_loss(teacher_immediate_layer.detach().float(), F.interpolate(immediate_layer.float(), teacher_immediate_layer.shape[2:], mode='nearest'), reduction="mean") 

                    loss = (
                        mse_loss + args.lpips_scale * lpips_loss + args.gen_scale * adaptive_weight * gen_loss + args.sc_scale * sc_loss
                    )

                    # Gather the losses across all processes for logging (if we use distributed training).
                    avg_loss = accelerator.gather(loss.repeat(args.train_batch_size)).mean()
                        
                    accelerator.backward(loss)
                    if accelerator.sync_gradients:
                        accelerator.clip_grad_norm_(vae.parameters(), args.max_grad_norm)
                    optimizer.step()
                    lr_scheduler.step()
                    # optimizer.zero_grad()
            else:
                # Return discriminator loss
                with accelerator.accumulate(discriminator):
                    pred.detach_()
                    target.requires_grad_()
                    real = discriminator(target)
                    fake = discriminator(pred)
                    loss = (F.relu(1 + fake) + F.relu(1 - real)).mean()
                    gp = gradient_penalty(target, real)
                    loss += gp
                    avg_discr_loss = accelerator.gather(loss.repeat(args.train_batch_size)).mean()

                    accelerator.backward(loss)
                    if accelerator.sync_gradients:
                        accelerator.clip_grad_norm_(discriminator.parameters(), args.max_grad_norm)
                    discr_optimizer.step()
                    discr_lr_scheduler.step()
                    

            # Checks if the accelerator has performed an optimization step behind the scenes
            if accelerator.sync_gradients:
                if args.use_ema:
                    ema_vae.step(vae.parameters())
                
                # teacher_vae.restore(vae.parameters())
                progress_bar.update(1)
                global_step += 1
            
            if accelerator.sync_gradients and not generator_step:
                logs = {
                    "gen_loss": gen_loss.detach().item(),
                    "lr": lr_scheduler.get_last_lr()[0],
                    "mse": mse_loss.detach().item(),
                    "lpips": lpips_loss.detach().item(),
                    # "kl": kl_loss.detach().item(),
                    "adaptive_weight": adaptive_weight.detach().item(),
                    "step_discr_loss": avg_discr_loss.detach().item(),
                    "sc_loss": sc_loss.detach().item(),
                }
                if avg_loss is not None:
                    logs["step_loss"] = avg_loss.detach().item(),

                accelerator.log(logs)
                progress_bar.set_postfix(**logs)

            # if accelerator.sync_gradients and global_step % num_update_steps_per_epoch == 0:
            if accelerator.sync_gradients and global_step % args.checkpointing_steps == 0:
                accelerator.wait_for_everyone()
                
                save_path = os.path.join(
                    args.output_dir, f"checkpoint-{global_step}"
                )
                accelerator.save_state(save_path)
                logger.info(f"Saved state to {save_path}")
            
            # Stop training if max steps is reached
            if global_step >= args.max_train_steps:
                break
            

        if accelerator.is_main_process:
            if epoch % args.validation_epochs == 0:
                if args.use_ema:
                    # Store the parameters temporarily and load the EMA parameters to perform inference.
                    ema_vae.store(vae.parameters())
                    ema_vae.copy_to(vae.parameters())
                with torch.no_grad():
                    log_validation(args, test_dataloader, vae, accelerator, weight_dtype, epoch, prefix="VAE F16")
                    if epoch == 0:
                        log_validation(args, test_dataloader, teacher_vae, accelerator, weight_dtype, epoch, prefix="VAE F8")
                if args.use_ema:
                    # Switch back to the original parameters.
                    ema_vae.restore(vae.parameters())
        
        

    # Create the pipeline using the trained modules and save it.
    accelerator.wait_for_everyone()
    if accelerator.is_main_process:
        vae = accelerator.unwrap_model(vae)
        discriminator = accelerator.unwrap_model(discriminator)
        if args.use_ema:
            ema_vae.copy_to(vae.parameters())
        vae.save_pretrained(args.output_dir)
        discriminator.save_pretrained(os.path.join(args.output_dir, "discriminator"))

    accelerator.end_training()


if __name__ == "__main__":
    main()